import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial import Delaunay
from shapely.geometry import Polygon, MultiPolygon, MultiPoint, MultiLineString
from shapely.ops import polygonize, unary_union


def alpha_shape(points, alpha):
    if len(points) < 4:
        return MultiPoint(list(points)).convex_hull

    tri = Delaunay(points)
    edges = set()

    for ia, ib, ic in tri.simplices:
        pa, pb, pc = points[ia], points[ib], points[ic]
        a = np.linalg.norm(pa - pb)
        b = np.linalg.norm(pb - pc)
        c = np.linalg.norm(pc - pa)
        s = (a + b + c) / 2.0
        area = max(s * (s - a) * (s - b) * (s - c), 1e-10)
        circ_radius = a * b * c / (4.0 * np.sqrt(area))
        if circ_radius < 1.0 / alpha:
            edges.add((ia, ib))
            edges.add((ib, ic))
            edges.add((ic, ia))

    edge_points = [(points[edge[0]], points[edge[1]]) for edge in edges]
    m = MultiLineString(edge_points)
    triangles = list(polygonize(m))
    shape = unary_union(triangles)

    if isinstance(shape, MultiPolygon):
        shape = max(shape.geoms, key=lambda p: p.area)
    elif not isinstance(shape, Polygon):
        shape = MultiPoint(list(points)).convex_hull

    return shape


def plot_alpha_shape(ax, points, alpha, color, label=None):
    shape = alpha_shape(points, alpha)
    if isinstance(shape, Polygon):
        x, y = shape.exterior.xy
        ax.plot(x, y, linestyle="--", color=color, linewidth=1, alpha=0.8, label=label)


def plot_tsne_embeddings_with_classes(file_paths, labels_paths, num_classes, save_path=None):
    fig = plt.figure(figsize=(18, 10))
    colors = ['r', 'g', 'b', 'y', 'm', 'c', 'k', 'orange', 'purple', 'pink', 'brown', 'gray', 'navy', 'teal', 'olive']
    handles = []
    methods = ['VIT','OVIT','HEAD']
    # methods = ['OVIT','HEAD']
    # methods = ['VIT','HEAD']
    # 第一行：完整的嵌入可视化
    for i, (file_path, label_path) in enumerate(zip(file_paths, labels_paths)):
        ax = fig.add_axes([0.08 + i * 0.3, 0.55, 0.3, 0.4])  # [left, bottom, width, height]
        embeddings_2d = np.load(file_path)
        labels = np.load(label_path)

        unique_classes = np.unique(labels)
        # unique_classes = [2,3,5,6]
        for class_index in unique_classes:
            class_mask = labels[:, 0] == class_index
            class_embeddings = embeddings_2d[class_mask]
            scatter = ax.scatter(class_embeddings[:, 0], class_embeddings[:, 1],color=colors[class_index], s=3,edgecolors='none', marker='o', label=f'Class {class_index}')
            # scatter = ax.scatter(
            #     class_embeddings[:, 0],
            #     class_embeddings[:, 1],
            #     color=colors[class_index],
            #     s=2,  # 增大点的大小
            #     marker='o',
            #     # alpha=0.3,  # 增加透明度
            #     edgecolor='k',  # 添加黑色边框，提升对比度
            #     label=f'Class {class_index}'
            # )
            # ax.grid(True, linestyle='--', alpha=0.5)  # 添加网格线
            if i == 0:  # 只添加一次图例
                handles.append(scatter)

        ax.axis('off')
        ax.text(0.5, -0.05, f'{methods[i]}', ha='center', va='center', transform=ax.transAxes, fontsize=12)
    handles2= []
    # 第二行：每类 alpha shape
    for class_index in unique_classes:
        ax = fig.add_axes([0.05 + class_index * 0.14, 0.2, 0.12, 0.25])  # [left, bottom, width, height]
        for i, (file_path, label_path) in enumerate(zip(file_paths, labels_paths)):
            embeddings_2d = np.load(file_path)
            labels = np.load(label_path)

            class_mask = labels[:, 0] == class_index
            class_embeddings = embeddings_2d[class_mask]
            plot_alpha_shape(ax, class_embeddings, alpha=0.6, color=colors[i])

        ax.axis('off')
        ax.text(0.5, -0.05, f'Class {class_index}', ha='center', va='center', transform=ax.transAxes, fontsize=12)

    # 在第二行子图的上方添加共享图例
    legend_ax = fig.add_axes([0.2, 0.48, 0.6, 0.03])  # [left, bottom, width, height]
    # legend_ax.axis('off')
    # legend_ax.legend(handles2, [f'Class {i}' for i in range(3)],
    #                  loc='center', ncol=3, fontsize=12, frameon=False)
    legend_ax.axis('off')
    legend_handles = [plt.Line2D([0], [0], color=colors[i], linestyle='--', lw=1, label=f'File {i + 1}')
                      for i in range(len(file_paths))]
    legend_ax.legend(legend_handles, [f'{methods[i]}' for i in range(len(file_paths))],
                     loc='center', ncol=len(file_paths), fontsize=12, frameon=False)

    # 保存图像
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0)


# 示例调用
file_paths = [
    "plots/tsne/ISIC/frozen_full/ISIC_embeddings_2d_vit.npy",
    "plots/tsne/ISIC/frozen_full/ISIC_embeddings_2d_ovit.npy",
    "plots/tsne/ISIC/head/ISIC_embeddings_2d_vit.npy",
]

labels_paths = [
    "plots/tsne/ISIC/frozen_full/ISIC_labels_vit.npy",
    "plots/tsne/ISIC/frozen_full/ISIC_labels_ovit.npy",
    "plots/tsne/ISIC/head/ISIC_labels_vit.npy"
]

plot_tsne_embeddings_with_classes(file_paths, labels_paths, num_classes=7, save_path="plots/tsne/ISIC/combined_embeddings_classes2.png")
